package org.example.test.brave;

import brave.Span;
import brave.Tag;
import brave.Tracer;
import brave.Tracing;
import brave.baggage.BaggageField;
import brave.baggage.BaggagePropagation;
import brave.baggage.BaggagePropagationConfig;
import brave.propagation.B3Propagation;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import org.apache.http.HttpException;
import org.apache.http.HttpRequest;
import org.apache.http.HttpRequestInterceptor;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.protocol.HttpContext;
import org.apache.http.util.EntityUtils;
import zipkin2.reporter.brave.AsyncZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import java.io.IOException;

public class LearnBrave {

    public static OkHttpSender sender = OkHttpSender.create("http://127.0.0.1:9411/api/v2/spans");
    public static AsyncZipkinSpanHandler zipkinSpanHandler = AsyncZipkinSpanHandler.create(sender);
    /* BaggageField */
    public static BaggageField COUNTRY_CODE = BaggageField.create("country-code");

    public static Tracing tracing = Tracing.newBuilder()
            .localServiceName("tester")
            .addSpanHandler(zipkinSpanHandler)
            .propagationFactory(
                    BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
                            .add(BaggagePropagationConfig.SingleBaggageField.remote(COUNTRY_CODE))
                            .build())
            .build();

    public static void client() {
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();

        Propagation.RemoteSetter<HttpUriRequest> SETTER = new Propagation.RemoteSetter<HttpUriRequest>() {
            @Override
            public Span.Kind spanKind() {
                return Span.Kind.CLIENT;
            }

            @Override
            public void put(HttpUriRequest request, String name, String value) {
                request.setHeader(name, value);
            }
        };

        TraceContext.Injector<HttpUriRequest> injector = tracing.propagation().injector(SETTER);

        Tracer tracer = tracing.tracer();

        /* In-process Tracing */
        // Start a new trace or a span within an existing trace representing operation
//        ScopedSpan span = tracer.startScopedSpan("encode");
        Span span = tracer.nextSpan().name("client").start();
        // Put the span in "scope" so that downstream code such as loggers can see trace IDs
        try (/*Tracer.SpanInScope ws = tracer.withSpanInScope(span)*/
                CurrentTraceContext.Scope ws = currentTraceContext.newScope(span.context())) {
//            request();

            CloseableHttpClient httpClient = HttpClients.custom().build();
            HttpGet httpGet = new HttpGet("http://localhost:8081/provide");

            injector.inject(span.context(), httpGet); // 注入上下文
            span.kind(Span.Kind.CLIENT);
            span.start();
            CloseableHttpResponse response = httpClient.execute(httpGet);
            System.out.println("Response: " + response.getStatusLine());
            System.out.println(EntityUtils.toString(response.getEntity()));
            response.close();
            httpClient.close();

            /* Customizing spans */
            span.tag("clnt/finagle.version", "6.36.0");

            Tag<String> TEST_TAG = new Tag<String>("test-tag") {
                @Override
                protected String parseValue(String input, TraceContext context) {
                    return input;
                }
            };
            TEST_TAG.tag("qinbiao", span);

        } catch (Exception | Error e) {
            span.error(e);
            e.printStackTrace();
        } finally {
            span.finish();
        }

        try {
            System.in.read();
        } catch (IOException e) {
            e.printStackTrace();
        }

        tracing.close();
        zipkinSpanHandler.close();
        sender.close();
    }

    public static void main(String[] args) {
        client();
    }

    public static void request() throws IOException {
        CloseableHttpClient httpClient = HttpClients.custom().addInterceptorLast(new HttpRequestInterceptor() {
            @Override
            public void process(HttpRequest request, HttpContext context) throws HttpException, IOException {
                request.setHeader("X-B3-TraceId", "xxx");
                request.setHeader("X-B3-ParentSpanId", "xxx");
                request.setHeader("X-B3-SpanId", "xxx");
                request.setHeader("X-B3-Sampled", "xxx");
            }
        }).build();
        HttpGet httpGet = new HttpGet("http://localhost:8081/test");
        CloseableHttpResponse response = httpClient.execute(httpGet);
        System.out.println("Response: " + response.getStatusLine());
        System.out.println(EntityUtils.toString(response.getEntity()));
        response.close();
        httpClient.close();
    }
}
